# 导入相关模块
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
from torch import nn
from torch import optim
from sklearn.datasets import make_regression
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def create_dataset():
    x, y, conf = make_regression(n_samples=100,
                                 n_features=1,
                                 random_state=0,
                                 noise=10,
                                 bias=10,
                                 coef=True)

    x = torch.tensor(x, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)
    conf = torch.tensor(conf, dtype=torch.float32)
    return x, y, conf


if __name__ == '__main__':
    x, y, conf = create_dataset()

    # 构造数据集对象
    dataset = TensorDataset(x, y)
    data_loader = DataLoader(dataset, batch_size=10, shuffle=True)

    # 构造线性回归模型
    model = nn.Linear(in_features=1, out_features=1)
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(params=model.parameters(), lr=0.001)

    epochs = 100
    loss_epoch = []

    for epoch in range(epochs):
        total_loss = 0.0
        for train_x, train_y in data_loader:
            y_pred = model(train_x)
            loss_values = loss_fn(y_pred, train_y.reshape(-1, 1))

            optimizer.zero_grad()
            loss_values.backward()
            optimizer.step()

            total_loss += loss_values.item()

        avg_loss = total_loss / len(data_loader)
        loss_epoch.append(avg_loss)

    # 绘制损失值变化曲线
    loss_epoch_np = torch.tensor(loss_epoch).detach().numpy()  # 转换为numpy数组
    plt.plot(range(1, epochs + 1), loss_epoch_np)
    plt.title('训练损失变化')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.show()

    # 绘制预测结果
    plt.scatter(x, y, label='真实数据')
    x1 = torch.linspace(x.min(), x.max(), 1000)
    y0 = model(x1).detach()  # 分离梯度
    y1 = x1 * conf + 1.5

    plt.plot(x1, y0, label='预测模型')
    plt.plot(x1, y1, label='真实模型')
    plt.grid()
    plt.legend()
    plt.show()